from mmengine.config import read_base
with read_base():
    from .mgam import *

from mmengine.model.weight_init import PretrainedInit
from mgamdata.models.SegFormer3D import (
    SegFormer3D_Encoder_MM, 
    SegFormer3D_Decoder_MM)
from mgamdata.mm.mmseg_Dev3D import DiceLoss_3D, EncoderDecoder_3D


embed_dims = 32
embed_dims_SegFormer3D = [
    embed_dims, 2*embed_dims, 5*embed_dims, 8*embed_dims]
head_embed_dims = embed_dims_SegFormer3D[-1]//2
ckpt_path = "/fileser51/zhangyiqin.sx/mmseg/FastSlow/work_dirs/1.1.3.1.UseMomentum/SegFormer3D/iter_50000.pth"
ckpt_path = None

# 神经网络设定
model = dict(
    type = EncoderDecoder_3D,
    backbone=dict(
        type=SegFormer3D_Encoder_MM, 
        init_cfg = dict(
            type=PretrainedInit, 
            checkpoint=ckpt_path, 
            prefix="backbone.0", 
            map_location="cpu"
        ) if ckpt_path is not None else None,
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims_SegFormer3D, 
        freeze=False,
    ), 
    decode_head=dict(
        type=SegFormer3D_Decoder_MM, 
        init_cfg = dict(
            type=PretrainedInit, 
            checkpoint=ckpt_path,
            prefix="backbone.1",
            map_location="cpu"
        ) if ckpt_path is not None else None,
        embed_dims=embed_dims_SegFormer3D, 
        head_embed_dims=head_embed_dims, 
        num_classes=num_classes, 
        loss_gt_key='gt_sem_seg', # ["gt_sem_seg_one_hot", "gt_sem_seg"]
        loss_decode=dict(
            type=DiceLoss_3D, 
            use_sigmoid=False, 
            ignore_1st_index=False, 
            batch_z=4, 
            # NOTE Severe performance overhead when not being set to None.
            # NOTE Prefer using `ignore_1st_index`.
            # NOTE Invalid Class (Defaults to the last class) has been masked out during preprocess.
            ignore_index=None, 
        ), 
    ),
    test_cfg=dict(
        mode="slide",
        crop_size=size,
        slide_accumulate_device='cpu',
        stride=[i//2 for i in size]
    )
)
